#!/usr/bin/env python3
"""
简化的进度管理模块
整合了所有进度管理功能，替代原来的分散实现
"""

import os
import time
import json
import argparse
import re

def extract_stance(text, split_by="操控陈述"):
    """从文本中提取指定标记的内容部分"""
    # 根据split_by动态构建多种可能的分割标记
    possible_splits = [
        f"【{split_by}】：",
        f"【{split_by}】",
        f"**{split_by}**",
        f"**{split_by}：**"
    ]
    
    # 尝试找到匹配的分割标记
    found_split = None
    for split_marker in possible_splits:
        if split_marker in text:
            found_split = split_marker
            break
    
    if found_split:
        stance_text = text.split(found_split, 1)[1].strip()
        
        # 使用正则表达式找到下一个【】或**括起来的部分
        next_split = re.search(r'【.*?】', stance_text)
        if next_split:
            return stance_text[:next_split.start()].strip()
        
        return stance_text  
    return ""

def save_progress_with_results(current_index, total_count, results, progress_file="progress.json"):
    """保存当前进度和已完成的实验结果"""
    progress_data = {
        "current_index": current_index,
        "total_count": total_count,
        "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
        "results": results  # 保存已完成的实验结果
    }
    
    try:
        with open(progress_file, "w", encoding="utf-8") as f:
            json.dump(progress_data, f, ensure_ascii=False, indent=2)
        print(f"💾 进度已保存: {current_index}/{total_count} (包含 {len(results)} 个结果)")
        return True
    except Exception as e:
        print(f"❌ 保存进度失败: {e}")
        return False

def load_progress_with_results(progress_file="progress.json"):
    """加载进度和已完成的实验结果"""
    if not os.path.exists(progress_file):
        print(f"📋 未找到进度文件: {progress_file}")
        return 0, []
    
    try:
        with open(progress_file, "r", encoding="utf-8") as f:
            progress_data = json.load(f)
        
        current_index = progress_data.get("current_index", 0)
        total_count = progress_data.get("total_count", 0)
        timestamp = progress_data.get("timestamp", "未知")
        results = progress_data.get("results", [])
        
        print(f"🔄 从进度恢复: {current_index}/{total_count}")
        print(f"   保存时间: {timestamp}")
        print(f"   已完成的实验数: {len(results)}")
        
        return current_index, results
    except Exception as e:
        print(f"❌ 加载进度失败: {e}")
        return 0, []

def cleanup_progress(progress_file="progress.json"):
    """清理进度文件"""
    if os.path.exists(progress_file):
        try:
            os.remove(progress_file)
            print(f"🧹 进度文件已清理: {progress_file}")
            return True
        except Exception as e:
            print(f"❌ 清理进度失败: {e}")
            return False
    return False

def should_save_progress(current_index, total_count, save_interval=5):
    """判断是否应该保存进度"""
    return (current_index + 1) % save_interval == 0 or (current_index + 1) == total_count

def backup_results(results, backup_file=None):
    """备份实验结果"""
    if not backup_file:
        backup_file = f"results_backup_{int(time.time())}.json"
    
    try:
        os.makedirs("results", exist_ok=True)
        backup_path = os.path.join("results", backup_file)
        
        with open(backup_path, "w", encoding="utf-8") as f:
            json.dump(results, f, ensure_ascii=False, indent=2)
        
        print(f"💾 结果已备份: {backup_path}")
        return True
    except Exception as e:
        print(f"❌ 备份结果失败: {e}")
        return False

def parse_progress_arguments():
    """解析进度管理相关的命令行参数"""
    parser = argparse.ArgumentParser(description="人格操控实验 - 支持进度恢复")
    parser.add_argument("--no-resume", action="store_true", 
                       help="不从进度恢复，重新开始实验")
    parser.add_argument("--progress-file", type=str, default="progress.json",
                       help="指定进度文件名")
    parser.add_argument("--save-interval", type=int, default=5,
                       help="进度保存间隔（每完成多少个实验保存一次）")
    parser.add_argument("--no-cleanup", action="store_true",
                       help="实验完成后不清理进度文件")
    parser.add_argument("--backup-results", action="store_true",
                       help="备份已完成的实验结果")
    
    return parser.parse_args()

def get_progress_config(args=None):
    """获取进度管理配置"""
    if args is None:
        # 使用默认配置
        return {
            "enabled": True,  # 是否启用进度管理
            "save_interval": 5,  # 每完成多少个实验保存一次进度
            "progress_file": "progress.json",  # 进度文件名
            "auto_cleanup": True,  # 实验完成后是否自动清理进度文件
            "resume_from_progress": True,  # 是否从进度恢复
            "backup_results": False,  # 是否备份结果
        }
    else:
        # 使用命令行参数
        return {
            "enabled": True,
            "save_interval": args.save_interval,
            "progress_file": args.progress_file,
            "auto_cleanup": not args.no_cleanup,
            "resume_from_progress": not args.no_resume,
            "backup_results": args.backup_results,
        }

# 向后兼容的简化函数
def save_progress(current_index, total_count, progress_file="progress.json"):
    """保存当前进度（向后兼容）"""
    return save_progress_with_results(current_index, total_count, [], progress_file)

def load_progress(progress_file="progress.json"):
    """加载进度（向后兼容）"""
    current_index, _ = load_progress_with_results(progress_file)
    return current_index

# 测试函数
def test_progress_manager():
    """测试进度管理功能"""
    
    # 模拟实验结果
    test_results = [
        {"id": "1", "status": "completed", "data": "test1"},
        {"id": "2", "status": "completed", "data": "test2"},
        {"id": "3", "status": "completed", "data": "test3"}
    ]
    
    progress_file = "test_progress_manager.json"
    
    # 测试保存进度
    print("=== 测试保存进度 ===")
    save_progress_with_results(3, 10, test_results, progress_file)
    
    # 测试加载进度
    print("\n=== 测试加载进度 ===")
    current_index, loaded_results = load_progress_with_results(progress_file)
    print(f"当前索引: {current_index}")
    print(f"加载的结果数: {len(loaded_results)}")
    
    # 测试备份结果
    print("\n=== 测试备份结果 ===")
    backup_results(loaded_results)
    
    # 清理测试文件
    cleanup_progress(progress_file)
    print("\n=== 测试完成 ===")

if __name__ == "__main__":
    test_progress_manager() 